package ffte.types

import breeze.math.Complex
import spinal.core._
import spinal.lib._

import ffte.property.{getResolution,getComplexMulMethod}

/**
 * Fixed-point complex number implementation companion object
 *
 * This object provides factory methods and utility functions for creating and
 * manipulating fixed-point complex numbers in hardware. It includes optimized
 * summation algorithms with configurable pipeline depths and precision management.
 *
 * Key Features:
 * - Vector-to-complex conversion utilities
 * - Pipelined summation with configurable depth
 * - Precision-aware operations with automatic bit growth management
 * - Recursive summation for optimal resource usage
 *
 * Hardware Optimization:
 * - Pipeline depth control for timing optimization
 * - Bit-width management to prevent overflow
 * - Resource-efficient recursive algorithms
 * - Configurable complex multiplication methods
 */
object FixComplex {
    /**
     * Convert 2-element vector to complex number
     *
     * Converts a Vec[FixReal] of length 2 to a FixComplex number.
     * The first element becomes the real part, the second becomes imaginary.
     *
     * @param x Input vector with 2 real elements [real, imaginary]
     * @return FixComplex number with converted components
     * @throws AssertionError if vector length is not 2
     */
    def fromVec(x:Vec[FixReal]) : FixComplex = {
        assert(x.length==2,s"only vec 2 real can convert to complex")
        val r = FixComplex(x(0).width,x(0).resolution)
        r.re := x(0)
        r.im := x(1)
        r
    }
  /**
     * Pipelined summation of complex numbers with configurable depth
     *
     * Implements a resource-efficient summation algorithm that supports configurable
     * pipeline depth for timing optimization. Uses recursive divide-and-conquer
     * approach for balanced tree structure when pipeline depth allows.
     *
     * Algorithm Strategy:
     * - d == 1: Single-cycle combinational sum with register output
     * - d > log2Up(sz): Use combinational sum with delay registers
     * - d < log2Up(sz): Use recursive pipelined summation tree
     *
     * @param a Vector of complex numbers to sum
     * @param sz Size of the vector (should equal a.length)
     * @param shift Right shift amount for precision management
     * @param d Pipeline depth (number of cycles for complete operation)
     * @param ce Clock enable signal for pipeline control
     * @return Summed complex number with managed precision
     */
    def sum(a:Vec[FixComplex],sz:Int,shift:Int,d:Int,ce:Bool):FixComplex = {
        val sW  = log2Up(sz)
        val dW  = a(0).width
        val res = a(0).resolution
        if(d==1) {
            val r = RegInit(FixComplex(dW,res).zero)
            val suma = new Area {
                val b = a.map(_.resize(dW+sW))
                when(ce) {
                    if(sz==1) {
                        r := b(0).cut(dW,shift)
                    } else {
                        r := b.reduce(_+_).fixTo(dW,shift)
                    }
                }
            }
            r
        }
        else if(d>sW) Delay(sum(a,sz,shift,sW,ce),d-sW,ce)
        else {
            val r = RegInit(FixComplex(dW,res).zero)
            val suma = new Area {
                when(ce) {
                    if(sz==1) {
                        r := Delay(a(0).resize(dW+sW).cut(dW,shift),d-1,ce)
                    } else {
                        val half = sz/2
                        val other = sz-half
                        val s = if(shift>=1) 1 else 0
                        r := (sum(Vec((0 until half).map(a(_))),half,shift-s,d-1,ce) + sum(Vec((half until sz).map(a(_))),other,shift-s,d-1,ce)).fixTo(dW,s)
                    }
                }
            }
            r
        }
    }
  /**
     * Pipelined summation with precision cutting
     *
     * Similar to sum() but applies precision cutting (cut instead of fixTo)
     * in the final combination stage. This provides alternative precision
     * management behavior for specific hardware optimization needs.
     *
     * Difference from sum():
     * - Uses cut() instead of fixTo() in the recursive combination
     * - Provides different rounding/saturation behavior
     * - Useful for applications requiring specific precision characteristics
     *
     * @param a Vector of complex numbers to sum
     * @param sz Size of the vector
     * @param shift Right shift amount for precision management
     * @param d Pipeline depth
     * @param ce Clock enable signal
     * @return Summed complex number with precision cutting
     */
    def sumc(a:Vec[FixComplex],sz:Int,shift:Int,d:Int,ce:Bool):FixComplex = {
        val sW = log2Up(sz)
        val dW = a(0).width
        val res = a(0).resolution
        if(d==1) {
            val r = RegInit(FixComplex(dW,res).zero)
            val suma = new Area {
                val b = a.map(_.resize(dW+sW))
                when(ce) {
                    if(sz==1) {
                        r := b(0).cut(dW,shift)
                    } else {
                        r := b.reduce(_+_).cut(dW,shift)
                    }
                }
            }
            r
        }
        else if(d>sW) Delay(sum(a,sz,shift,sW,ce),d-sW,ce)
        else {
            val r = RegInit(FixComplex(dW,res).zero)
            val suma = new Area {
                when(ce) {
                    if(sz==1) {
                        r := Delay(a(0).resize(dW+sW).cut(dW,shift),d-1,ce)
                    } else {
                        val half = sz/2
                        val other = sz-half
                        val s = if(shift>=1) 1 else 0
                        r := (sum(Vec((0 until half).map(a(_))),half,shift-s,d-1,ce) + sum(Vec((half until sz).map(a(_))),other,shift-s,d-1,ce)).cut(dW,s)
                    }
                }
            }
            r
        }
    }
}

/**
 * Fixed-point complex number implementation
 *
 * This class implements a fixed-point complex number as a SpinalHDL Bundle with
 * separate real and imaginary components. It provides comprehensive arithmetic
 * operations, precision management, and utility functions optimized for hardware
 * implementation in FFT systems.
 *
 * Data Representation:
 * - Complex number with real (re) and imaginary (im) components
 * - Configurable bit width and fractional resolution
 * - Signed two's complement representation
 * - Hardware-optimized arithmetic operations
 *
 * Key Features:
 * - Complex arithmetic operations (+, -, *, unary minus)
 * - Precision management (resize, fixTo, cut, sat)
 * - Special operations (conjugate, absolute, rotation by j)
 * - Mathematical functions (exponential, trigonometric)
 * - Conversion utilities (to/from software Complex, vector operations)
 *
 * Hardware Optimizations:
 * - Configurable complex multiplication methods
 * - Automatic bit growth management
 * - Resource-efficient implementations
 * - Pipeline-friendly operations
 *
 * @param width Bit width for both real and imaginary components
 * @param resolution Fractional bit resolution (defaults to global setting)
 */
case class FixComplex(val width:Int,val resolution:Int=getResolution()) extends Bundle {
    /** Real component of the complex number */
    val re = FixReal(width,resolution)

    /** Imaginary component of the complex number */
    val im = FixReal(width,resolution)
    /**
     * Expand bit width by 1 bit
     *
     * Creates a new FixComplex with increased bit width for both components.
     * Useful for preventing overflow in addition operations.
     *
     * @return New FixComplex with width+1 bits per component
     */
    def expand = {
        val ret = FixComplex(width+1,resolution)
        ret.re := this.re.expand
        ret.im := this.im.expand
        ret
    }
    /**
     * Complex addition
     *
     * Performs component-wise addition of two complex numbers.
     * Output bit width matches input width (no automatic growth).
     *
     * @param right Complex number to add
     * @return Sum of this and right complex numbers
     */
    def + (right: FixComplex): FixComplex = {
        val ret = FixComplex(width,resolution)
        ret.re := this.re + right.re
        ret.im := this.im + right.im
        ret
    }

    /**
     * Complex subtraction
     *
     * Performs component-wise subtraction of two complex numbers.
     * Output bit width matches input width (no automatic growth).
     *
     * @param right Complex number to subtract
     * @return Difference of this and right complex numbers
     */
    def - (right: FixComplex): FixComplex = {
        val ret = FixComplex(width,resolution)
        ret.re := this.re - right.re
        ret.im := this.im - right.im
        ret
    }
    
    /**
     * Assignment operator
     *
     * Assigns the values from another FixComplex to this complex number.
     * Used for signal assignment in SpinalHDL.
     *
     * @param that Source complex number to assign from
     */
    def :=(that: FixComplex): Unit     = {
        this.re := that.re
        this.im := that.im
    }

    /**
     * Complex multiplication
     *
     * Performs complex multiplication using configurable methods for optimization.
     * Supports both direct and reduced-multiplication algorithms.
     *
     * Mathematical Formula:
     * (a + bj) * (c + dj) = (ac - bd) + (ad + bc)j
     *
     * Implementation Options:
     * - Method 1 (getComplexMulMethod() = true): Direct 4-multiplication approach
     * - Method 2: Reduced 3-multiplication approach using algebraic optimization
     *
     * @param right Complex multiplier
     * @return Product with expanded bit width and resolution
     */
    def * (that: FixComplex): FixComplex = {
        val ret = FixComplex(this.width+that.width,this.resolution+that.resolution)
        if(getComplexMulMethod()) {
            // Direct 4-multiplication method: (ac - bd) + (ad + bc)j
            ret.re.d.raw := this.re.d.raw*that.re.d.raw - this.im.d.raw*that.im.d.raw
            ret.im.d.raw := this.im.d.raw*that.re.d.raw + this.re.d.raw*that.im.d.raw
        } else{
            // Reduced 3-multiplication method using algebraic identity:
            // Let p = (a+b)(c+d) = ac + ad + bc + bd
            // Then real = ac - bd, imag = p - ac - bd = ad + bc
            val this_add = this.re.resize(width+1) + this.im.resize(width+1)
            val that_add = that.re.resize(width+1) + that.im.resize(width+1)
            val re_re    = this.re * that.re
            val im_im    = this.im * that.im
            val add_add  = this_add * that_add
            ret.re := re_re - im_im
            ret.im := (add_add - re_re - im_im).cut(this.width+that.width,0)
        }
        ret
    }

    /**
     * Complex-real multiplication
     *
     * Multiplies a complex number by a real number, scaling both components.
     * More efficient than complex-complex multiplication.
     *
     * @param that Real multiplier
     * @return Scaled complex number with expanded bit width
     */
    def * (that: FixReal): FixComplex = {
        val ret = FixComplex(this.width+that.width)
        ret.re.d.raw := this.re.d.raw*that.d.raw
        ret.im.d.raw := this.im.d.raw*that.d.raw
        ret
    }
    /**
     * Unary negation
     *
     * Negates both real and imaginary components.
     * Returns -(re + im*j) = -re - im*j
     *
     * @return Negated complex number
     */
    def unary_- : FixComplex = {
        val ret = FixComplex(width,resolution)
        ret.re := -this.re
        ret.im := -this.im
        ret
    }
    /**
     * Multiply by imaginary unit j
     *
     * Rotates the complex number by 90 degrees counter-clockwise.
     * Returns j * (re + im*j) = -im + re*j
     *
     * @return Complex number rotated by 90 degrees
     */
    def j = {
        val ret = FixComplex(width,resolution)
        ret.re := -this.im
        ret.im := this.re
        ret
    }
    /**
     * Complex conjugate
     *
     * Returns the complex conjugate by negating the imaginary part.
     * Returns (re + im*j)* = re - im*j
     *
     * @return Complex conjugate
     */
    def conj = {
        val ret = FixComplex(width,resolution)
        ret.re := this.re
        ret.im := -this.im
        ret
    }
    /**
     * Complex absolute value approximation
     *
     * Computes an efficient approximation of |re + im*j| using the
     * α-max-plus-beta-min algorithm with α=1, β=0.5.
     * More hardware-efficient than the exact sqrt(re² + im²).
     *
     * Approximation: |a + bj| ≈ max(|a|, |b|) + 0.5*min(|a|, |b|)
     *
     * @return Approximate magnitude as unsigned integer
     */
    def cabs : UInt = {
        val rabs = this.re.abs
        val iabs = this.im.abs
        (rabs>iabs) ? (rabs+(iabs |>> 1)) | ((rabs |>> 1)+iabs)
    }
    def fixTo(w:Int,s:Int) = {
        val ret = FixComplex(w,resolution)
        ret.re := this.re.fixTo(w,s)
        ret.im := this.im.fixTo(w,s)
        ret
    }
    def fixTo(w:Int,s:Int,r:Int) = {
        val ret = FixComplex(w,r)
        ret.re := this.re.fixTo(w,s,r)
        ret.im := this.im.fixTo(w,s,r)
        ret
    }
    def cut(w:Int,s:Int) = {
        val ret = FixComplex(w,resolution)
        ret.re := this.re.cut(w,s)
        ret.im := this.im.cut(w,s)
        ret
    }
    def cut(w:Int,s:Int,r:Int) = {
        val ret = FixComplex(w,r)
        ret.re := this.re.cut(w,s,r)
        ret.im := this.im.cut(w,s,r)
        ret
    }
    def one = {
        val ret = FixComplex(width,resolution)
        ret.re := this.re.one
        ret.im := this.im.zero
        ret
    }
    def zero = {
        val ret = FixComplex(width,resolution)
        ret.re.d.raw := S(0)
        ret.im.d.raw := S(0)
        ret
    }
    def resize(dW:Int) = {
        val r = FixComplex(dW,resolution)
        r.re := this.re.resize(dW)
        r.im := this.im.resize(dW)
        r
    }
    def sat(m:Int) = {
        val r = FixComplex(width)
        r.re := this.re.sat(m)
        r.im := this.im.sat(m)
        r
    }
    def expj(m:Int,N:Int) = {
        val r = FixComplex(width,resolution) 
        r.re := this.re.cos(m,N)
        r.im := this.re.sin(m,N)
        r
    }
    def swap = {
        val r = FixComplex(width,resolution)
        r.re := this.im
        r.im := this.re
        r
    }
    def toVec = Vec(this.re,this.im)
    def fromComplex(x:Complex) = {
        this.re.fromDouble(x.re)
        this.im.fromDouble(x.im)
        this
    }
}

